﻿/*  Copyright (C) 2012 AmaroK86 (marcidm 'at' hotmail 'dot' com)
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.

 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.

 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
using System;
using System.IO;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Gibbed.IO;
using AmaroK86.MassEffect3.ZlibBlock;

namespace AmaroK86.MassEffect3.PCC
{
    public class PCCObject
    {
        public string pccFileName { get; private set; }

        static int headerSize = 0x8E;
        byte[] header = new byte[headerSize];
        //byte[] extraNamesList = null;

        private uint   magic    { get { return BitConverter.ToUInt32(header, 0); } }
        private ushort lowVers  { get { return BitConverter.ToUInt16(header, 4); } }
        private ushort highVers { get { return BitConverter.ToUInt16(header, 6); } }
        private int    nameSize { get { int val = BitConverter.ToInt32(header, 12); if (val < 0) return val * -2; else return val; } }
        public uint    flags    { get { return BitConverter.ToUInt32(header, 16 + nameSize); } }

        public bool isModified { get { return Exports.Any(entry => entry.hasChanged == true); } }
        public bool bDLCStored = false;
        public bool bExtraNamesList { get { return (flags & 0x10000000) != 0; } }
        public bool bCompressed
        {
            get { return (flags & 0x02000000) != 0; }
            set
            {
                if (value) // sets the compressed flag if bCompressed set equal to true
                    Buffer.BlockCopy(BitConverter.GetBytes(flags | 0x02000000), 0, header, 16 + nameSize, sizeof(int));
                else // else set to false
                    Buffer.BlockCopy(BitConverter.GetBytes(flags & ~0x02000000), 0, header, 16 + nameSize, sizeof(int));
            }
        }

        int idxOffsets   { get { if ((flags & 8) != 0) return 24 + nameSize; else return 20 + nameSize; } }
        int NameCount    { get { return BitConverter.ToInt32(header, idxOffsets); }      set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets,      sizeof(int)); } }
        int NameOffset   { get { return BitConverter.ToInt32(header, idxOffsets + 4); }  set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets + 4,  sizeof(int)); } }
        int ExportCount  { get { return BitConverter.ToInt32(header, idxOffsets + 8); }  set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets + 8,  sizeof(int)); } }
        int ExportOffset { get { return BitConverter.ToInt32(header, idxOffsets + 12); } set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets + 12, sizeof(int)); } }
        int ImportCount  { get { return BitConverter.ToInt32(header, idxOffsets + 16); } set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets + 16, sizeof(int)); } }
        int ImportOffset { get { return BitConverter.ToInt32(header, idxOffsets + 20); } set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets + 20, sizeof(int)); } }

        int expInfoEndOffset { get { return BitConverter.ToInt32(header, idxOffsets + 24); } set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets + 24, sizeof(int)); } }
        int expDataBegOffset { get { return BitConverter.ToInt32(header, idxOffsets + 28); } set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, idxOffsets + 28, sizeof(int)); Buffer.BlockCopy(BitConverter.GetBytes(value), 0, header, 8, sizeof(int)); } }

        int headerEnd;

        public List<string> Names { get; private set; }
        public ImportList Imports { get; private set; }
        public ExportList Exports { get; private set; }

        List<Block> blockList = null;
        protected class Block
        {
            public int uncOffset;
            public int uncSize;
            public int cprOffset;
            public int cprSize;
            public bool bRead = false;
        }

        public class ImportEntry
        {
            public static int byteSize = 28;
            internal byte[] data = new byte[byteSize];
            internal PCCObject pccRef;
            public int Link;

            public int idxPackageFile { get { return BitConverter.ToInt32(data, 0); }      private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, data, 0, sizeof(int)); } }
            public int idxClassName   { get { return BitConverter.ToInt32(data, 8); }      private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, data, 8, sizeof(int)); } }
            public int idxPackageName { get { return BitConverter.ToInt32(data, 16) - 1; } private set { Buffer.BlockCopy(BitConverter.GetBytes(value + 1), 0, data, 16, sizeof(int)); } }
            public int idxObjectName  { get { return BitConverter.ToInt32(data, 20); }     private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, data, 20, sizeof(int)); } }
            public int idxLink        { get { return BitConverter.ToInt32(data, 16); }     private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, data, 16, sizeof(int)); } }
            public int ObjectFlags    { get { return BitConverter.ToInt32(data, 24); }     private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, data, 24, sizeof(int)); } }

            public string ClassName   { get { return pccRef.Names[idxClassName]; } }
            public string PackageFile { get { return pccRef.Names[idxPackageFile] + ".pcc"; } }
            public string ObjectName  { get { return pccRef.Names[idxObjectName]; } }
            public string PackageName { get { int val = idxPackageName; if (val >= 0) return pccRef.Names[pccRef.Exports[val].idxObjectName]; else return "Package"; } }
            public string PackageFullName
            {
                get
                {
                    string result = PackageName;
                    int idxNewPackName = idxPackageName;

                    while (idxNewPackName >= 0)
                    {
                        string newPackageName = pccRef.Exports[idxNewPackName].PackageName;
                        if (newPackageName != "Package")
                            result = newPackageName + "." + result;
                        idxNewPackName = pccRef.Exports[idxNewPackName].idxPackageName;
                    }
                    return result;
                }
            }

            public ImportEntry(PCCObject pccFile, byte[] importData)
            {
                pccRef = pccFile;
                data = (byte[])importData.Clone();
            }

            public ImportEntry(PCCObject pccFile, Stream importData)
            {
                pccRef = pccFile;
                data = new byte[ImportEntry.byteSize];
                importData.Read(data, 0, data.Length);
            }
        }

        public class ImportList : List<ImportEntry>
        {
            // this indexing allows easier access to data editing
            public ImportEntry this[string searchName]
            {
                get
                {
                    int idx = this.FindIndex(import => import.ObjectName == searchName);
                    if (idx != -1)
                        return this[idx];
                    else
                        throw new KeyNotFoundException("Import name \"" + searchName + "\" not found");
                }
                set
                {
                    int idx = this.FindIndex(import => import.ObjectName == searchName);
                    if (idx != -1)
                        this[idx] = value;
                    else
                        throw new KeyNotFoundException("Import name \"" + searchName + "\" not found");
                }
            }
        }

        public class ExportEntry : ICloneable // class containing info about export entry (header info + data)
        {
            internal byte[] info; // holds data about export header, not the export data.
            public PCCObject pccRef;
            public uint offset { get; private set; }
            public int Link; // deprecated var, soon will be removed

            public int idxClassName   { get { return BitConverter.ToInt32(info, 0); }     private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, info, 0, sizeof(int)); } }
            public int idxClassParent { get { return BitConverter.ToInt32(info, 4); }     private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, info, 4, sizeof(int)); } }
            public int idxPackageName { get { return BitConverter.ToInt32(info, 8) - 1; } private set { Buffer.BlockCopy(BitConverter.GetBytes(value + 1), 0, info, 8, sizeof(int)); } }
            public int idxObjectName  { get { return BitConverter.ToInt32(info, 12); }    private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, info, 12, sizeof(int)); } }
            //public int archetypeIdx   { get { return BitConverter.ToInt32(info, 20); } private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, info, 20, sizeof(int)); } }
            public long ObjectFlags   { get { return BitConverter.ToInt64(info, 24); }    private set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, info, 64, sizeof(long)); } }

            public string ObjectName  { get { return pccRef.Names[idxObjectName]; } }
            public string ClassName   { get { int val = idxClassName;   if (val < 0)  return pccRef.Names[pccRef.Imports[val * -1 - 1].idxObjectName]; else if (val > 0) return pccRef.Names[pccRef.Exports[val].idxObjectName]; else return "Class"; } }
            public string ClassParent { get { int val = idxClassParent; if (val < 0)  return pccRef.Names[pccRef.Imports[val * -1 - 1].idxObjectName]; else if (val > 0) return pccRef.Names[pccRef.Exports[val].idxObjectName]; else return "Class"; } }
            public string PackageName { get { int val = idxPackageName; if (val >= 0) return pccRef.Names[pccRef.Exports[val].idxObjectName]; else return "Package"; } }
            public string PackageFullName
            {
                get
                {
                    string result = PackageName;
                    int idxNewPackName = idxPackageName;

                    while (idxNewPackName >= 0)
                    {
                        string newPackageName = pccRef.Exports[idxNewPackName].PackageName;
                        if (newPackageName != "Package")
                            result = newPackageName + "." + result;
                        idxNewPackName = pccRef.Exports[idxNewPackName].idxPackageName;
                    }
                    return result;
                }
            }

            internal int originalDataSize { get { return BitConverter.ToInt32(info, 32); } set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, info, 32, sizeof(int)); } }
            public int DataOffset { get { return BitConverter.ToInt32(info, 36); } internal set { Buffer.BlockCopy(BitConverter.GetBytes(value), 0, info, 36, sizeof(int)); } }
            byte[] _data = null;
            public byte[] Data // holds data about export data
            {
                get
                {
                    // if data isn't loaded then fill it from pcc file (load-on-demand)
                    if (_data == null)
                        pccRef.getData(DataOffset);
                    return _data;
                }

                set { _data = value; hasChanged = true; }
            }
            public bool hasChanged { get; internal set; }

            public ExportEntry(PCCObject pccFile, byte[] importData, uint exportOffset)
            {
                pccRef = pccFile;
                info = (byte[])importData.Clone();
                offset = exportOffset;
                hasChanged = false;
            }

            object ICloneable.Clone()
            {
                return this.Clone();
            }

            public ExportEntry Clone()
            {
                ExportEntry newExport = (ExportEntry)this.MemberwiseClone(); // copy all reference-types vars
                // now creates new copies of referenced objects
                newExport.info = (byte[])this.info.Clone();
                newExport.Data = (byte[])this.Data.Clone();
                return newExport;
            }
        }

        public class ExportList : List<ExportEntry>
        {
            // this indexing allows easier access to data editing
            public ExportEntry this[string searchName, string className = null]
            {
                get
                {
                    int idx = (className == null)
                        ? this.FindIndex(export => export.ObjectName == searchName)
                        : this.FindIndex(export => export.ObjectName == searchName && export.ClassName == className);
                    if (idx != -1)
                        return this[idx];
                    else
                        throw new KeyNotFoundException("Export name \"" + searchName + "\" not found");
                }
                set
                {
                    int idx = this.FindIndex(export => export.ObjectName == searchName);
                    if (idx != -1)
                        this[idx] = value;
                    else
                        throw new KeyNotFoundException("Export name \"" + searchName + "\" not found");
                }
            }
        }

        /// <summary>
        ///     PCCObject class constructor. It also load namelist, importlist and exportinfo (not exportdata) from pcc file
        /// </summary>
        /// <param name="pccFilePath">full path + file name of desired pcc file.</param>
        public PCCObject(string pccFilePath, bool fullFileInMemory = false)
        {
            pccFileName = Path.GetFullPath(pccFilePath);
            using (FileStream pccStream = File.OpenRead(pccFileName))
            {
                Names = new List<string>();
                Imports = new ImportList();
                Exports = new ExportList();

                pccStream.Read(header, 0, header.Length);
                if (magic != ZBlock.magic &&
                    magic.Swap() != ZBlock.magic)
                {
                    throw new FormatException("not a pcc file");
                }

                if (lowVers != 684 && highVers != 194)
                {
                    throw new FormatException("unsupported version");
                }

                Stream listsStream;

                if (bCompressed)
                {
                    // seeks the blocks info position
                    pccStream.Seek(idxOffsets + 60, SeekOrigin.Begin);
                    int generator = pccStream.ReadValueS32();
                    pccStream.Seek((generator * 12) + 20, SeekOrigin.Current);

                    int blockCount = pccStream.ReadValueS32();
                    blockList = new List<Block>();

                    // creating the Block list
                    for (int i = 0; i < blockCount; i++)
                    {
                        Block temp = new Block();
                        temp.uncOffset = pccStream.ReadValueS32();
                        temp.uncSize = pccStream.ReadValueS32();
                        temp.cprOffset = pccStream.ReadValueS32();
                        temp.cprSize = pccStream.ReadValueS32();
                        blockList.Add(temp);
                    }

                    // correcting the header, in case there's need to be saved
                    Buffer.BlockCopy(BitConverter.GetBytes((int)0), 0, header, header.Length - 12, sizeof(int));
                    pccStream.Read(header, header.Length - 8, 8);
                    headerEnd = (int)pccStream.Position;

                    // decompress first block that holds infos about names, imports and exports
                    pccStream.Seek(blockList[0].cprOffset, SeekOrigin.Begin);
                    byte[] uncBlock = ZBlock.Decompress(pccStream, blockList[0].cprSize);

                    // write decompressed block inside temporary stream
                    listsStream = new MemoryStream();
                    listsStream.Seek(blockList[0].uncOffset, SeekOrigin.Begin);
                    listsStream.Write(uncBlock, 0, uncBlock.Length);
                }
                else
                {
                    listsStream = pccStream;
                    headerEnd = (int)NameOffset;
                }

                /*if(bExtraNamesList)
                {
                    int extraNamesListSize = namesOffset - headerEnd;
                    extraNamesList = new byte[extraNamesListSize];
                    pccStream.Seek(headerEnd, SeekOrigin.Begin);
                    pccStream.Read(extraNamesList, 0, extraNamesList.Length);
                }*/

                // fill names list
                listsStream.Seek(NameOffset, SeekOrigin.Begin);
                for (int i = 0; i < NameCount; i++)
                {
                    int strLength = listsStream.ReadValueS32();
                    Names.Add(listsStream.ReadString(strLength * -2, true, Encoding.Unicode));
                }

                // fill import list
                listsStream.Seek(ImportOffset, SeekOrigin.Begin);
                byte[] buffer = new byte[ImportEntry.byteSize];
                for (int i = 0; i < ImportCount; i++)
                {
                    Imports.Add(new ImportEntry(this, listsStream));
                }

                // fill export list (only the headers, not the data)
                listsStream.Seek(ExportOffset, SeekOrigin.Begin);
                for (int i = 0; i < ExportCount; i++)
                {
                    uint expInfoOffset = (uint)listsStream.Position;

                    listsStream.Seek(44, SeekOrigin.Current);
                    int count = listsStream.ReadValueS32();
                    listsStream.Seek(-48, SeekOrigin.Current);

                    int expInfoSize = 68 + (count * 4);
                    buffer = new byte[expInfoSize];

                    listsStream.Read(buffer, 0, buffer.Length);
                    Exports.Add(new ExportEntry(this, buffer, expInfoOffset));
                }
            }
        }

        /// <summary>
        ///     given export data offset, the function recovers it from the file.
        /// </summary>
        /// <param name="offset">offset position of desired export data</param>
        private void getData(int offset)
        {
            byte[] buffer;
            if (bCompressed)
            {
                Block selected = blockList.Find(block => block.uncOffset <= offset && block.uncOffset + block.uncSize > offset);
                byte[] uncBlock;

                using (FileStream pccStream = File.OpenRead(pccFileName))
                {
                    pccStream.Seek(selected.cprOffset, SeekOrigin.Begin);
                    uncBlock = ZBlock.Decompress(pccStream, selected.cprSize);

                    // the selected block has been read
                    selected.bRead = true;
                }

                // fill all the exports data extracted from the uncBlock
                foreach (ExportEntry expInfo in Exports)
                {
                    if (expInfo.DataOffset >= selected.uncOffset && expInfo.DataOffset + expInfo.originalDataSize <= selected.uncOffset + selected.uncSize)
                    {
                        buffer = new byte[expInfo.originalDataSize];
                        Buffer.BlockCopy(uncBlock, expInfo.DataOffset - selected.uncOffset, buffer, 0, expInfo.originalDataSize);
                        expInfo.Data = buffer;
                    }
                }
            }
            else
            {
                int expIndex = Exports.FindIndex(export => export.DataOffset <= offset && export.DataOffset + export.originalDataSize > offset);
                ExportEntry expSelect = Exports[expIndex];
                using (FileStream pccStream = File.OpenRead(pccFileName))
                {
                    buffer = new byte[expSelect.originalDataSize];
                    pccStream.Seek(expSelect.DataOffset, SeekOrigin.Begin);
                    pccStream.Read(buffer, 0, buffer.Length);
                    expSelect.Data = buffer;
                }
            }
        }

        /// <summary>
        ///     save PCCObject to original file.
        /// </summary>
        /// <param name="saveCompress">set true if you want a zlib compressed pcc file.</param>
        public void saveToFile(bool saveCompress)
        {
            saveToFile(null, saveCompress);
        }

        
        /// <summary>
        ///     save PCCObject to file.
        /// </summary>
        /// <param name="newFileName">set full path + file name.</param>
        /// <param name="saveCompress">set true if you want a zlib compressed pcc file.</param>
        public void saveToFile(string newFileName = null, bool saveCompress = false)
        {
            bool bOverwriteFile = false;

            if (string.IsNullOrWhiteSpace(newFileName) || newFileName == pccFileName)
            {
                bOverwriteFile = true;
                newFileName = Path.GetFullPath(pccFileName) + ".tmp";
            }

            if (bDLCStored)
                saveCompress = false;

            using (MemoryStream newPccStream = new MemoryStream())
            {
                // this condition needs a deeper check. todo...
                /*if (bExtraNamesList)
                {
                    newPccStream.Seek(headerEnd, SeekOrigin.Begin);
                    newPccStream.Write(extraNamesList, 0, extraNamesList.Length);
                }*/

                //writing names list
                newPccStream.Seek(NameOffset, SeekOrigin.Begin);
                NameCount = Names.Count;
                foreach (string name in Names)
                {
                    newPccStream.WriteValueS32(-(name.Length + 1));
                    newPccStream.WriteString(name + "\0", (uint)(name.Length + 1) * 2, Encoding.Unicode);
                }

                //writing import infos
                ImportOffset = (int)newPccStream.Position;
                ImportCount = Imports.Count;
                foreach (ImportEntry import in Imports)
                    newPccStream.Write(import.data, 0, import.data.Length);

                //updating general export infos
                ExportOffset = (int)newPccStream.Position;
                ExportCount = Exports.Count;
                expInfoEndOffset = ExportOffset + Exports.Sum(export => export.info.Length);
                if(expDataBegOffset < expInfoEndOffset)
                    expDataBegOffset = expInfoEndOffset;

                //writing export data
                newPccStream.Seek(expDataBegOffset, SeekOrigin.Begin);
                foreach (ExportEntry export in Exports)
                {
                    //updating info values
                    export.originalDataSize = export.Data.Length;
                    export.DataOffset = (int)newPccStream.Position;

                    //writing data
                    newPccStream.Write(export.Data, 0, export.Data.Length);
                }

                //writing export info
                newPccStream.Seek(ExportOffset, SeekOrigin.Begin);
                foreach (ExportEntry export in Exports)
                {
                    newPccStream.Write(export.info, 0, export.info.Length);
                }

                //writing header
                bCompressed = false;
                newPccStream.Seek(0, SeekOrigin.Begin);
                newPccStream.Write(header, 0, header.Length);
                newPccStream.Seek(0, SeekOrigin.Begin);

                if (saveCompress)
                {
                    PCCHandler.CompressAndSave(newPccStream, newFileName);
                }
                else
                {
                    using (FileStream outputStream = File.Create(newFileName))
                    {
                        newPccStream.CopyTo(outputStream);
                    }
                }
            }

            if (bOverwriteFile)
            {
                File.Delete(pccFileName);
                File.Move(newFileName, pccFileName);
            }
        }

        public void addName(string name)
        {
            Names.Add(name);
        }

        public void addImport(PCCObject.ImportEntry importEntry)
        {
            if (importEntry.pccRef != this)
                throw new Exception("you cannot add a new import entry from another pcc file, it has invalid references!");

            Imports.Add(importEntry);
        }

        public void addExport(PCCObject.ExportEntry exportEntry)
        {
            if (exportEntry.pccRef != this)
                throw new Exception("you cannot add a new export entry from another pcc file, it has invalid references!");

            exportEntry.hasChanged = true;

            //changing data offset in order to append it at the end of the file
            ExportEntry lastExport = Exports.Find(export => export.DataOffset == Exports.Max(entry => entry.DataOffset));
            int lastOffset = lastExport.DataOffset + lastExport.Data.Length;
            exportEntry.DataOffset = lastOffset;

            Exports.Add(exportEntry);
        }
    }
}